import torch
import torch.nn as nn

def compare_models(model1, model2, atol=1e-6):
    """
    比较两个模型是否相同，包含结构、权重、输出。
    
    参数:
        model1: 第一个模型
        model2: 第二个模型
        atol: 比较浮动精度的容差值，默认1e-6
        
    返回:
        None
    """
    # 1. 比较模型结构
    print("比较模型结构...")
    if str(model1) == str(model2):
        print("模型结构相同!")
    else:
        print("模型结构不同!")

    # 2. 比较模型参数
    print("\n比较模型参数...")
    state_dict1 = model1.state_dict()
    state_dict2 = model2.state_dict()

    # 比较每个参数
    for key in state_dict1:
        if key not in state_dict2:
            print(f"参数 {key} 在第二个模型中缺失")
        else:
            if torch.allclose(state_dict1[key], state_dict2[key], atol=atol):
                print(f"参数 {key} 在两个模型中相同")
            else:
                print(f"参数 {key} 在两个模型中不同")

    # 3. 比较模型输出
    print("\n比较模型输出...")
    input_data = torch.randn(1, 3, 224, 224)  # 示例输入（根据需要调整）
    
    output1 = model1(input_data)
    output2 = model2(input_data)

    if torch.allclose(output1, output2, atol=atol):
        print("两个模型的输出相同!")
    else:
        print("两个模型的输出不同!")


# 示例：假设有两个模型 model1 和 model2
# model1 = <你自己的模型1>
# model2 = <你自己的模型2>

# 运行比较
# compare_models(SwinTansformerUnetPP, SwinTransformerSys)
